package com.yj.timecapsule.websocket.ai;
import cn.hutool.json.JSONUtil;
import com.yj.timecapsule.common.enums.AiMessageType;
import com.yj.timecapsule.common.enums.ErrorCode;
import com.yj.timecapsule.common.utils.AsyncTaskUtil;
import com.yj.timecapsule.mapper.AiMessageMapper;
import com.yj.timecapsule.service.ai.ChatAiService;
import com.yj.timecapsule.websocket.ai.entity.AiMessage;
import com.yj.timecapsule.websocket.ai.entity.MessageContent;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import java.io.IOException;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
/**
 * @author Sean
 * @CreateTime 2024-08-05  11:09
 * @Description
 * @Version 1.0
 */
@Component
@Slf4j
public class WebSocketCallHandler extends TextWebSocketHandler {
    @Resource
    private AiMessageMapper aiMessageMapper;
    private static final int MAX_TOTAL = 100;
    @Resource
    private ChatAiService chatAiService;
    @Resource
    private ThreadPoolTaskExecutor websocketExecutor;
    private static final List<AiMessage> HISTORY = Collections.synchronizedList(new ArrayList<>());
    private static final int MAX_CONNECTION = 500;
    private static final Map<Long,WebSocketSession> SESSIONS = new ConcurrentHashMap<>();
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        String payload = message.getPayload();
        String content = chatAiService.call(payload);
        log.info("当前线程:{}", Thread.currentThread().getName());
        log.info("内容:{}", content);
        long id = Long.parseLong(session.getAttributes().get("userId").toString());
        saveHistoryMessage(id,payload, content);
        AsyncTaskUtil.runAsync(()->{
            log.info("当前线程:{}", Thread.currentThread().getName());
            try {
                Random random = new Random();
                StringBuilder dataBuilder = new StringBuilder();
                for (char c : content.toCharArray()) {
                    dataBuilder.append(c);
                }
                int totalLength = dataBuilder.length();
                int sentLength = 0;
                while (sentLength < totalLength) {
                    int randomNumber = random.nextInt(6) + 1;  // 生成 1 到 6 之间的随机数
                    int toSendLength = Math.min(randomNumber, totalLength - sentLength);
                    String toSend = dataBuilder.substring(sentLength, sentLength + toSendLength);
                    session.sendMessage(new TextMessage(toSend));
                    sentLength += toSendLength;
                    Thread.sleep(150);
                }
            } catch (Exception e) {
                try {
                    session.close();
                } catch (IOException ex) {
                    log.error("关闭失败", ex);
                }
            }
        },websocketExecutor).whenComplete((result, throwable) -> {
            if (throwable != null){
                try {
                    session.sendMessage(new TextMessage("连接出现异常，请重试"));
                } catch (IOException e) {
                    try {
                        session.close();
                    } catch (IOException ex) {
                        throw new RuntimeException(ex);
                    }
                }
            }
        });

    }

    private void saveHistoryMessage(Long userId,String payload, String content) {
        AiMessage aiMessage = new AiMessage();
        aiMessage.setUserId(userId);
        aiMessage.setCreateTime(new Date());
        MessageContent userMessageContent = new MessageContent();
        userMessageContent.setType(AiMessageType.USER.getValue());
        userMessageContent.setContent(payload);
        MessageContent aiMessageContent = new MessageContent();
        aiMessageContent.setType(AiMessageType.AI.getValue());
        aiMessageContent.setContent(content);
        List<MessageContent> list = new ArrayList<>();
        list.add(userMessageContent);
        list.add(aiMessageContent);
        String jsonStr = JSONUtil.toJsonStr(list);
        aiMessage.setContent(jsonStr);
        HISTORY.add(aiMessage);
        if (HISTORY.size() == MAX_TOTAL){
            aiMessageMapper.insertBatch(HISTORY);
            HISTORY.clear();
        }
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        if (SESSIONS.size() >= MAX_CONNECTION){
            session.sendMessage(new TextMessage(ErrorCode.CONNECTION_TO_MANY.getMsg()));
            session.close();
            return;
        }
        Long userId = Long.parseLong(session.getAttributes().get("userId").toString());
        SESSIONS.put(userId,session);
        log.info("建立连接 当前全部连接={}",SESSIONS);

    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        Long userId = Long.parseLong(session.getAttributes().get("userId").toString());
        SESSIONS.remove(userId);
        log.info("关闭连接 userid={}",userId);
        super.afterConnectionClosed(session, status);
    }
}
